{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8744d20",
   "metadata": {},
   "source": [
    "### LLM Inference Example\n",
    "\n",
    "This notebook contains a basic inference example for using our `ttml` Python API to build, load, and run a large language model from Hugging Face on our TT hardware. By default, it is set to create and load a GPT2 model, but this notebook can quickly and easily be edited to use any of the LLMs that the tt-train project currently supports. \n",
    "\n",
    "Below, in the first cell, we have our imports and basic directory housekeeping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a7a20c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, random\n",
    "import numpy as np  # For numpy arrays\n",
    "from dataclasses import dataclass # For configuration classes\n",
    "from huggingface_hub import hf_hub_download # To download safetensors from Hugging Face\n",
    "from transformers import AutoTokenizer\n",
    "from yaml import safe_load # To read YAML configs\n",
    "from pathlib import Path\n",
    "\n",
    "sys.path.append(f\"{os.environ['TT_METAL_HOME']}/tt-train/sources/ttml\")\n",
    "import ttml\n",
    "from ttml.common.config import get_config, TransformerConfig\n",
    "from ttml.common.utils import set_seed, round_up_to_tile\n",
    "from ttml.common.model_factory import TransformerModelFactory\n",
    "\n",
    "# Change working directory to tt-train\n",
    "os.chdir(f\"{os.environ['TT_METAL_HOME']}/tt-train\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e005a45c",
   "metadata": {},
   "source": [
    "Use the cell below to change global parameters in this notebook. \n",
    "\n",
    "`OUTPUT_TOKENS` : the length of the generated text in token (not characters!) \n",
    "\n",
    "`WITH_SAMPLING` : enable or disable output token sampling (only used for PyTorch)\n",
    "\n",
    "`TEMPERATURE`   : sampling temperature; set to 0 to disable sampling in `generate_with_tt()`\n",
    "\n",
    "`SEED`          : randomization seed (for reproducibility)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a87d6dc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_TOKENS = 256\n",
    "WITH_SAMPLING = True\n",
    "TEMPERATURE = 0.8\n",
    "SEED = 42\n",
    "CONFIG = \"gpt2_inference.yaml\"\n",
    "\n",
    "set_seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e9b032f",
   "metadata": {},
   "source": [
    "While the notebook is currently configured for GPT2, you can quickly change the tokenizer you want to use by changing the input to `from_pretrained()` below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ea71a9df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the tokenizer from Hugging Face and the transformer config from YAML\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "# transformer_config = TransformerConfig(get_config(CONFIG).get(\"training_config\", {}).get(\"transformer_config\",{}))\n",
    "yaml_config = get_config(CONFIG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5415d24d",
   "metadata": {},
   "source": [
    "As above, the call to `hf_hub_download()` will download (or otherwise find on your local system) the SafeTensors model weight file for GPT2, but can be updated to download other SafeTensors files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6965256c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Safetensors path: /home/ubuntu/.cache/huggingface/hub/models--gpt2/snapshots/607a30d783dfa663caf39e06633721c8d4cfcd7e/\n"
     ]
    }
   ],
   "source": [
    "# # Get safetensors\n",
    "safetensors_path = hf_hub_download(repo_id=\"gpt2\", filename=\"model.safetensors\")\n",
    "safetensors_path = safetensors_path.replace(\"model.safetensors\",\"\")\n",
    "\n",
    "print(f\"Safetensors path: {safetensors_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "907e5933",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_causal_mask(T: int) -> ttml.autograd.Tensor:\n",
    "    # [1,1,T,T] float32 with 1s for allowed positions (i >= j), else 0\\n\",\n",
    "    m = np.tril(np.ones((T, T), dtype=np.float32))\n",
    "    return ttml.autograd.Tensor.from_numpy(m.reshape(1, 1, T, T), ttml.Layout.TILE, ttml.autograd.DataType.BFLOAT16)\n",
    "\n",
    "def build_logits_mask(vocab_size: int, padded_vocab_size: int) -> ttml.autograd.Tensor:\n",
    "    logits_mask = np.zeros((1, 1, 1, padded_vocab_size), dtype=np.float32)\n",
    "    logits_mask[:, :, :, vocab_size:] = 1e4\n",
    "    return ttml.autograd.Tensor.from_numpy(logits_mask, ttml.Layout.TILE, ttml.autograd.DataType.BFLOAT16)   # [1,1,1,T], bfloat16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd89028f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer configuration:\n",
      "    Vocab size: 50272\n",
      "    Max sequence length: 1024\n",
      "    Embedding dim: 768\n",
      "    Num heads: 12\n",
      "    Dropout probability: 0.2\n",
      "    Num blocks: 12\n",
      "    Positional embedding type: Trainable\n",
      "    Runner type: Default\n",
      "    Composite layernorm: false\n",
      "    Weight tying: Enabled\n",
      "2025-10-03 00:58:33.846 | info     |          Device | Opening user mode device driver (tt_cluster.cpp:189)\n",
      "2025-10-03 00:58:33.900 | info     |             UMD | Harvesting mask for chip 0 is 0x80 (NOC0: 0x80, simulated harvesting mask: 0x0). (cluster.cpp:402)\n",
      "2025-10-03 00:58:33.936 | warning  |             UMD | init_detect_tt_device_numanodes(): Could not determine NumaNodeSet for TT device (physical_device_id: 0 pci_bus_id: 0000:04:00.0) (cpuset_lib.cpp:565)\n",
      "2025-10-03 00:58:33.936 | warning  |             UMD | Could not find NumaNodeSet for TT Device (physical_device_id: 0 pci_bus_id: 0000:04:00.0) (cpuset_lib.cpp:173)\n",
      "2025-10-03 00:58:33.937 | warning  |             UMD | bind_area_memory_nodeset(): Unable to determine TT Device to NumaNode mapping for physical_device_id: 0. Skipping membind. (cpuset_lib.cpp:450)\n",
      "2025-10-03 00:58:33.937 | warning  |             UMD | ---- ttSiliconDevice::init_hugepage: bind_area_to_memory_nodeset() failed (physical_device_id: 0 ch: 0). Hugepage allocation is not on NumaNode matching TT Device. Side-Effect is decreased Device->Host perf (Issue #893). (sysmem_manager.cpp:212)\n",
      "2025-10-03 00:58:33.937 | info     |             UMD | Opening local chip ids/PCIe ids: {0}/[0] and remote chip ids {} (cluster.cpp:251)\n",
      "2025-10-03 00:58:33.937 | info     |             UMD | All devices in cluster running firmware version: 18.10.0 (cluster.cpp:231)\n",
      "2025-10-03 00:58:33.937 | info     |             UMD | IOMMU: disabled (cluster.cpp:173)\n",
      "2025-10-03 00:58:33.937 | info     |             UMD | KMD version: 2.4.0 (cluster.cpp:176)\n",
      "2025-10-03 00:58:33.937 | info     |             UMD | Software version 6.0.0, Ethernet FW version 7.0.0 (Device 0) (cluster.cpp:1061)\n",
      "2025-10-03 00:58:33.940 | info     |             UMD | Pinning pages for Hugepage: virtual address 0x7f9bc0000000 and size 0x40000000 pinned to physical address 0x200000000 (pci_device.cpp:551)\n",
      "2025-10-03 00:58:33.967 | info     |       Inspector | Inspector RPC server listening on 50051 (rpc_server_controller.cpp:102)\n",
      "2025-10-03 00:58:35.163 | info     |           Metal | DPRINT enabled on device 0, worker core (x=0,y=0) (virtual (x=18,y=18)). (dprint_server.cpp:715)\n",
      "2025-10-03 00:58:35.163 | info     |           Metal | DPRINT Server attached device 0 (dprint_server.cpp:762)\n",
      "Loading model from: /home/ubuntu/.cache/huggingface/hub/models--gpt2/snapshots/607a30d783dfa663caf39e06633721c8d4cfcd7e/model.safetensors\n",
      "parameter name: transformer/gpt_block_9/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_9/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_9/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_9/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_8/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_8/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_8/attention/qkv_linear/bias\n",
      "parameter name: transfPadding vocab size for tilization: original 50257 -> padded 50272\n",
      "ormer/gpt_block_7/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_7/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_7/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_7/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln1/gamma\n",
      "parameter name: transformer/gpt_block_6/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_0/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_4/ln1/beta\n",
      "parameter name: transformer/gpt_block_8/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_0/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_8/ln2/beta\n",
      "parameter name: transformer/gpt_block_3/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_9/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/ln1/gamma\n",
      "parameter name: transformer/gpt_block_11/ln2/beta\n",
      "parameter name: transformer/gpt_block_3/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_7/ln2/beta\n",
      "parameter name: transformer/gpt_block_11/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_7/ln1/beta\n",
      "parameter name: transformer/gpt_block_1/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_6/ln1/gamma\n",
      "parameter name: transformer/gpt_block_5/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_7/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_6/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_4/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/ln2/beta\n",
      "parameter name: transformer/gpt_block_8/ln1/gamma\n",
      "parameter name: transformer/gpt_block_6/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_4/ln1/gamma\n",
      "parameter name: transformer/gpt_block_5/ln1/beta\n",
      "parameter name: transformer/gpt_block_7/ln2/gamma\n",
      "parameter name: transformer/gpt_block_3/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_0/mlp/fc1/bias\n",
      "parameter name: transformer/ln_fc/beta\n",
      "parameter name: transformer/gpt_block_1/ln1/beta\n",
      "parameter name: transformer/gpt_block_4/ln2/beta\n",
      "parameter name: transformer/gpt_block_2/ln2/beta\n",
      "parameter name: transformer/gpt_block_1/ln2/beta\n",
      "parameter name: transformer/gpt_block_7/ln1/gamma\n",
      "parameter name: transformer/gpt_block_3/ln1/gamma\n",
      "parameter name: transformer/gpt_block_8/ln2/gamma\n",
      "parameter name: transformer/gpt_block_11/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_1/ln2/gamma\n",
      "parameter name: transformer/gpt_block_0/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/ln2/gamma\n",
      "parameter name: transformer/pos_emb/weight\n",
      "parameter name: transformer/gpt_block_5/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_11/ln1/beta\n",
      "parameter name: transformer/gpt_block_9/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln2/gamma\n",
      "parameter name: transformer/gpt_block_11/ln1/gamma\n",
      "parameter name: transformer/gpt_block_1/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_5/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/ln2/gamma\n",
      "parameter name: transformer/gpt_block_10/ln1/gamma\n",
      "parameter name: transformer/gpt_block_4/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_9/ln1/gamma\n",
      "parameter name: transformer/gpt_block_0/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_5/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_10/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_11/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/ln1/gamma\n",
      "parameter name: transformer/fc/weight\n",
      "parameter name: transformer/gpt_block_11/ln2/gamma\n",
      "parameter name: transformer/ln_fc/gamma\n",
      "parameter name: transformer/gpt_block_8/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_9/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_6/ln2/gamma\n",
      "parameter name: transformer/gpt_block_4/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/ln2/gamma\n",
      "parameter name: transformer/gpt_block_1/ln1/gamma\n",
      "parameter name: transformer/gpt_block_0/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_1/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_3/ln2/gamma\n",
      "parameter name: transformer/gpt_block_9/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/ln1/beta\n",
      "parameter name: transformer/gpt_block_1/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_1/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_1/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_4/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_8/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_11/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_1/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_1/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_10/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_11/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_8/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_2/ln1/beta\n",
      "parameter name: transformer/gpt_block_10/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_5/ln2/beta\n",
      "parameter name: transformer/gpt_block_11/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_10/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_10/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_11/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_3/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_11/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_5/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_2/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_2/attention/out_linear/weight\n",
      "parameter name: transformer/gpt_block_6/ln2/beta\n",
      "parameter name: transformer/gpt_block_2/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_0/ln1/beta\n",
      "parameter name: transformer/gpt_block_10/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_2/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_3/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_3/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_8/ln1/beta\n",
      "parameter name: transformer/gpt_block_3/mlp/fc1/weight\n",
      "parameter name: transformer/gpt_block_6/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_4/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_4/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_9/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_10/ln1/beta\n",
      "parameter name: transformer/gpt_block_5/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_5/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_5/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_6/attention/qkv_linear/weight\n",
      "parameter name: transformer/gpt_block_6/mlp/fc2/bias\n",
      "parameter name: transformer/gpt_block_8/mlp/fc1/bias\n",
      "parameter name: transformer/gpt_block_7/attention/qkv_linear/bias\n",
      "parameter name: transformer/gpt_block_3/ln2/beta\n",
      "parameter name: transformer/gpt_block_6/mlp/fc2/weight\n",
      "parameter name: transformer/gpt_block_7/attention/out_linear/bias\n",
      "parameter name: transformer/gpt_block_7/attention/out_linear/weight\n",
      "Loading tensor: h.0.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.0.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_0/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.0.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_0/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.0.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_0/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.0.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.0.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.0.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.0.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_0/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.1.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.1.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_1/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.1.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_1/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.1.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_1/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.1.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.1.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.1.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.1.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_1/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.10.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.10.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_10/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.10.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_10/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.10.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_10/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.10.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.10.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.10.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.10.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_10/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.11.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.11.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_11/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.11.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_11/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.11.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_11/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.11.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.11.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.11.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.11.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_11/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.2.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.2.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_2/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.2.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_2/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.2.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_2/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.2.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.2.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.2.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.2.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_2/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.3.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.3.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_3/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.3.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_3/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.3.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_3/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.3.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.3.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.3.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.3.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_3/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.4.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.4.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_4/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.4.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_4/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.4.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_4/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.4.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.4.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.4.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.4.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_4/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.5.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.5.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_5/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.5.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_5/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.5.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_5/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.5.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.5.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.5.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.5.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_5/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.6.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.6.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_6/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.6.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_6/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.6.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_6/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.6.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.6.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.6.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.6.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_6/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.7.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.7.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_7/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.7.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_7/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.7.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_7/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.7.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.7.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.7.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.7.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_7/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.8.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.8.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_8/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.8.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_8/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.8.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_8/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.8.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.8.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.8.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.8.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_8/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: h.9.attn.bias, shape:Shape([1, 1, 1024, 1024]), format: F32\n",
      "Loading tensor: h.9.attn.c_attn.bias, shape:Shape([2304]), format: F32\n",
      " Parameter transformer/gpt_block_9/attention/qkv_linear/bias, shape: Shape([1, 1, 1, 2304])\n",
      "Loading tensor: h.9.attn.c_attn.weight, shape:Shape([768, 2304]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_9/attention/qkv_linear/weight, shape: Shape([1, 1, 2304, 768])\n",
      "Loading tensor: h.9.attn.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/attention/out_linear/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.attn.c_proj.weight, shape:Shape([768, 768]), format: F32\n",
      "Transposing!\n",
      " Parameter transformer/gpt_block_9/attention/out_linear/weight, shape: Shape([1, 1, 768, 768])\n",
      "Loading tensor: h.9.ln_1.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln1/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_1.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln1/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_2.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln2/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.ln_2.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/ln2/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.mlp.c_fc.bias, shape:Shape([3072]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc1/bias, shape: Shape([1, 1, 1, 3072])\n",
      "Loading tensor: h.9.mlp.c_fc.weight, shape:Shape([768, 3072]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc1/weight, shape: Shape([1, 1, 3072, 768])\n",
      "Transposing!\n",
      "Loading tensor: h.9.mlp.c_proj.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc2/bias, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: h.9.mlp.c_proj.weight, shape:Shape([3072, 768]), format: F32\n",
      " Parameter transformer/gpt_block_9/mlp/fc2/weight, shape: Shape([1, 1, 768, 3072])\n",
      "Transposing!\n",
      "Loading tensor: ln_f.bias, shape:Shape([768]), format: F32\n",
      " Parameter transformer/ln_fc/beta, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: ln_f.weight, shape:Shape([768]), format: F32\n",
      " Parameter transformer/ln_fc/gamma, shape: Shape([1, 1, 1, 768])\n",
      "Loading tensor: wpe.weight, shape:Shape([1024, 768]), format: F32\n",
      " Parameter transformer/pos_emb/weight, shape: Shape([1, 1, 1024, 768])\n",
      "Loading tensor: wte.weight, shape:Shape([50257, 768]), format: F32\n",
      " Parameter transformer/fc/weight, shape: Shape([1, 1, 50272, 768])\n",
      "Original shape 50257, 768Transformed shape 50272, 768\n"
     ]
    }
   ],
   "source": [
    "orig_vocab_size = tokenizer.vocab_size\n",
    "\n",
    "tt_model_factory = TransformerModelFactory(yaml_config)\n",
    "tt_model_factory.transformer_config.vocab_size = orig_vocab_size\n",
    "\n",
    "max_sequence_length = tt_model_factory.transformer_config.max_sequence_length\n",
    "\n",
    "tt_model = tt_model_factory.create_model()\n",
    "tt_model.load_from_safetensors(safetensors_path)\n",
    "tt_model\n",
    "\n",
    "padded_vocab_size = round_up_to_tile(orig_vocab_size, 32)\n",
    "\n",
    "if orig_vocab_size != padded_vocab_size:\n",
    "    print(f\"Padding vocab size for tilization: original {orig_vocab_size} -> padded {padded_vocab_size}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40913e88",
   "metadata": {},
   "source": [
    "`generate_with_tt()` uses TT hardware acceleration to generate output from the chosen LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "18e6550e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_with_tt(model, prompt_tokens):\n",
    "\n",
    "    ttml.autograd.AutoContext.get_instance().set_gradient_mode(ttml.autograd.GradMode.DISABLED)\n",
    "    model.eval()\n",
    "\n",
    "    logits_mask_tensor = None\n",
    "\n",
    "    if padded_vocab_size != orig_vocab_size:\n",
    "        logits_mask_tensor = build_logits_mask(orig_vocab_size, padded_vocab_size)\n",
    "\n",
    "    causal_mask = build_causal_mask(max_sequence_length)  # [1,1,seq_len,seq_len], float32\n",
    "    padded_prompt_tokens = np.zeros((1, 1, 1, max_sequence_length), \n",
    "                                    dtype=np.uint32)\n",
    "\n",
    "    start_idx = 0\n",
    "\n",
    "    print(\"************************************\")\n",
    "    for token_idx in range(OUTPUT_TOKENS):\n",
    "\n",
    "        if len(prompt_tokens) > max_sequence_length:\n",
    "            start_idx = len(prompt_tokens) - max_sequence_length\n",
    "\n",
    "        # padded_prompt_tokens[0, 0, 0, :transformer_cfg[\"max_sequence_length\"]] = 0\n",
    "        padded_prompt_tokens[0, 0, 0, :len(prompt_tokens)] = prompt_tokens[start_idx:]\n",
    "        padded_prompt_tensor = ttml.autograd.Tensor.from_numpy(\n",
    "            padded_prompt_tokens,\n",
    "            ttml.Layout.ROW_MAJOR,\n",
    "            ttml.autograd.DataType.UINT32)  # [1,1,1, max_seq_len], uint32\n",
    "\n",
    "        logits = model(padded_prompt_tensor, causal_mask)  # out=[1,1,seq_len, vocab_size], bf16\n",
    "\n",
    "\n",
    "        next_token_tensor = ttml.ops.sample.sample_op(logits, TEMPERATURE, np.random.randint(low=1e7), logits_mask_tensor)  # out=[1,1,seq_len,1], uint32\n",
    "\n",
    "        next_token_idx = max_sequence_length - 1 if len(prompt_tokens) > max_sequence_length else len(prompt_tokens) - 1\n",
    "        next_token = next_token_tensor.to_numpy().flatten()[next_token_idx]\n",
    "\n",
    "        output = tokenizer.decode(next_token)\n",
    "\n",
    "        prompt_tokens.append(next_token)\n",
    "        print(output, end='', flush=True)\n",
    "\n",
    "    print(\"\\n************************************\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7824c0cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_with_pytorch(prompt_tokens):\n",
    "    import torch\n",
    "    from transformers import AutoModelForCausalLM\n",
    "\n",
    "    torch.manual_seed(SEED)\n",
    "\n",
    "    torch_model = AutoModelForCausalLM.from_pretrained(\"gpt2\", dtype=torch.bfloat16)\n",
    "    torch_model.eval()\n",
    "    print(\"************************************\")\n",
    "    with torch.no_grad():\n",
    "        outputs = torch_model.generate(\n",
    "            prompt_tokens,\n",
    "            max_new_tokens=OUTPUT_TOKENS,\n",
    "            do_sample=WITH_SAMPLING, # Enable sampling\n",
    "            temperature=TEMPERATURE,   # Temperature for sampling\n",
    "            num_beams=1 # Use multinomial sampling (standard sampling)\n",
    "        )\n",
    "    \n",
    "    generated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
    "\n",
    "    for t in generated_text:\n",
    "        print(t)\n",
    "        \n",
    "    print(\"\\n************************************\\n\\n\"),"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6a42273b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating with TT:\n",
      "************************************\n",
      "2025-10-03 00:59:26.195 | info     |            Test | Small moreh_layer_norm algorithm is selected. (moreh_layer_norm_program_factory.cpp:168)\n",
      " cats are a little smaller than dogs, and dogs are more muscular.\n",
      "\n",
      "\"The animals have a very wide range of sensory experiences, and the human brain is much smaller. It's a very different place to be.\n",
      "\n",
      "\"The difference between cats and dogs is: cats are a little smaller than dogs, and dogs are more muscular.\n",
      "\n",
      "\"The animals have a very wide range of sensory experiences, and the human brain is much smaller. It's a very different place to be.\n",
      "\n",
      "\"The animals have a very wide range of sensory experiences, and the human brain is much larger. It's a very different place to be.\"\n",
      "\n",
      "The first time we saw the dog, I was shocked to find a small, fluffy dog.\n",
      "\n",
      "\"I'm not a huge dog fan,\" said Ms Bevins, who works as a veterinarian at the University of Cambridge.\n",
      "\n",
      "\"But I do think it's a very important animal, and I think the human brain is very different. I think it's a very different place to be.\"\n",
      "\n",
      "The dogs in the study were all female, but the two males were all male.\n",
      "\n",
      "\"It is very surprising, and it's very interesting,\" Ms Bevins said.\n",
      "\n",
      "************************************\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt_str = \"The difference between cats and dogs is:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d98a32d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating with TT:\n",
      "************************************\n",
      " ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( ( (\n",
      " (\n",
      " ( (\n",
      "\n",
      "\n",
      " F F 35 (\n",
      " F\n",
      "\n",
      "\n",
      " (\n",
      "\n",
      "\n",
      "\n",
      " 35\n",
      " 35 F ( F<|endoftext|> 35 F F F (\n",
      " (\n",
      " ( ( F ( 35 F ( F F\n",
      " 35 F (<|endoftext|> 35 F F 35 A\n",
      " 35 F A<|endoftext|>\n",
      " F<|endoftext|> F 35<|endoftext|><|endoftext|> F F (\n",
      "\n",
      " F\n",
      " F F 35 F 35 ( The\n",
      " F 35 F F F 35 I F 35 F 35 F F F A F ( F F ( F F F The F F F A F F F F A 35 F F The The F<|endoftext|> F A F F 35<|endoftext|> A<|endoftext|> F F I A F The F F The F The The F The I<|endoftext|> The F A The F F The F\n",
      " I The F F The The F 35 The F F<|endoftext|> 35 The\n",
      " The F The F The H The F The H 35 F The\n",
      " 35 The The A The The F\n",
      " 35 The<|endoftext|> F The F The<|endoftext|>\n",
      "************************************\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt_str = \"Compared to spoons, forks are meant to:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b67723",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = \"Bees are similar to:\"\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT:\")\n",
    "generate_with_tt(tt_model, prompt_tokens.copy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efe3330e",
   "metadata": {},
   "source": [
    "Now try your own prompt!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8beb647",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = input(\"Enter your prompt: \")\n",
    "\n",
    "prompt_tokens = tokenizer.encode(prompt_str)\n",
    "print(\"Generating with TT model:\")\n",
    "generate_with_tt(tt_model, prompt_tokens)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
